package com.suns.demo;

import com.alibaba.fastjson.JSONObject;
import com.suns.demo.vo.User;
import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.remoting.common.RemotingHelper;

import java.io.UnsupportedEncodingException;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;

public class DemoTransMessage01 {

    /**
     * 用于暂存尚未确认的事务消息
     */
    private static ConcurrentHashMap<Message, User> ongoingData = new ConcurrentHashMap<>();

    /**
     * 用于暂存已回滚的消息
     */
    private static ConcurrentHashMap<Message, User> rollbackData = new ConcurrentHashMap<>();

    // 用于模拟目的
    private static Random random = new Random();

    public static void main(String[] args) throws MQClientException, InterruptedException {

        TransactionMQProducer txProducer = new TransactionMQProducer("tx_producer_01");

        txProducer.setNamesrvAddr("localhost:9876");

        // 设置事务监听器
        txProducer.setTransactionListener(new TransactionListener() {
            @Override
            public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {

                // 执行本地事务，如MySQL的增删改等
                System.out.println("使用【" + arg + "】执行本地事务");
                ongoingData.put(msg, (User) arg);

//                int index = random.nextInt(3);
//                switch (index) {
//                    case 0:
                        // 会发起回查
                        return LocalTransactionState.UNKNOW;
//                    case 1:
//                        return LocalTransactionState.ROLLBACK_MESSAGE;
//                    case 2:
//                        return LocalTransactionState.COMMIT_MESSAGE;
//                    default:
//                        return LocalTransactionState.UNKNOW;
//                }
            }

            @Override
            public LocalTransactionState checkLocalTransaction(MessageExt msg) {
                // 用于回查
                int status = random.nextInt(3);

                switch (status) {
                    case 0:
                        System.err.println(new String(msg.getBody()) + "---事务状态---未知");
                        return LocalTransactionState.UNKNOW;
                    case 1:
                        System.err.println(new String(msg.getBody()) + "---事务消息------提交");
                        // 事务消息提交，从ongoingData移除
                        ongoingData.remove(msg);
                        return LocalTransactionState.COMMIT_MESSAGE;
                    case 2:
                        System.err.println(new String(msg.getBody()) + "---事务消息---------回滚");
                        // 事务消息回滚，从ongoingData中移到rollbackData中。
                        rollbackData.put(msg, ongoingData.get(msg));
                        return LocalTransactionState.ROLLBACK_MESSAGE;
                    default:
                        rollbackData.put(msg, ongoingData.get(msg));
                        return LocalTransactionState.ROLLBACK_MESSAGE;
                }
            }
        });

        // 启动事务生产者
        txProducer.start();

        User user = new User(1, "zhangsan", "zs12321");

        JSONObject json = (JSONObject) JSONObject.toJSON(user);

        String[] tags = new String[]{"TagA", "TagB", "TagC", "TagD", "TagE"};

        for (int i = 0; i < 10; i++) {
            try {
                Message msg =
                        new Message("tp_user", tags[i % tags.length], "KEY" + i,
                                json.toJSONString().getBytes(RemotingHelper.DEFAULT_CHARSET));

                SendResult sendResult = txProducer.sendMessageInTransaction(msg, user);

                System.out.printf("%s%n", sendResult);

                Thread.sleep(10);

            } catch (MQClientException | UnsupportedEncodingException | InterruptedException e) {
                e.printStackTrace();
            }
        }

        for (int i = 0; i < 100000; i++) {
            Thread.sleep(1000);
        }

        txProducer.shutdown();

    }
}
